import os 
import uuid
from fastapi import FastAPI, File, UploadFile, Query
from fastapi.responses  import JSONResponse
from paddleocr import PaddleOCR
from fastapi.staticfiles  import StaticFiles
import requests

app = FastAPI()

# 配置参数
# APP_PORT = 5001
APP_PORT = 8000
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}

# Metal加速配置 
os.environ['PADDLE_ENABLE_METAL']  = '1'

os.makedirs("uploads",  exist_ok=True)
app.mount("/static",  StaticFiles(directory='uploads'), name="static")

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.',  1)[1].lower() in ALLOWED_EXTENSIONS
  
# 初始化OCR引擎（启用多进程加速）
ocr_engine = PaddleOCR(
    det_model_dir='./models/ch_PP-OCRv4_det_server_infer',
    rec_model_dir='./models/ch_PP-OCRv4_rec_server_infer',
    use_mp=True,          # 启用多进程 
    total_process=8,     # M1 Max建议8线程
    show_log=True,
    use_angle_cls=True,
    use_gpu=True,
    save_crop_images=False,
)

@app.get('/img-ocr')
async def img_to_ocr(url: str = Query(..., description="图片地址参数")):
    try:
        print(f'img url:{url}')
        result = ocr_engine.ocr(url,  cls=True)
        texts = [line[1][0] for res in result for line in res]
        print(f'ocr result:{texts}')

        # 删除tmp.jpg
        # if os.path.exists('tmp.jpg'):
        tmp_path = os.path.join('./', 'tmp.jpg')
        print(f'tmp_path:{tmp_path}')
        os.remove(tmp_path)

        return JSONResponse(texts)
    except Exception as e:
        print(f'error:{str(e)}')
        return JSONResponse([])

@app.get('/ocr')
async def ocr_service(file: UploadFile = File(...)):
    try:
        print(file)
        # # 检查文件是否存在
        if not file.filename  or file.filename  == 'undefined':
            return JSONResponse(
                status_code=400,
                content={'error': 'No file part'}
            )
        
        # file = request.files['file'] 
        
        # 检查是否选择了文件
        if file.filename  == '':
            return JSONResponse(
                status_code=400,
                content={'error': 'No selected file'}
            )
        
        # 验证文件类型
        if not allowed_file(file.filename): 
            return JSONResponse(
                status_code=400,
                content={'error': 'Invalid file type. Only support png, jpg, jpeg, gif'}
            )
        
        # 创建存储目录
        if not os.path.exists(UPLOAD_FOLDER): 
            os.makedirs(UPLOAD_FOLDER) 

        unique_name = f"{uuid.uuid4().hex}{file.content_type.replace('image/', '.')}"
        
        # 保存文件
        save_path = os.path.join(UPLOAD_FOLDER,  unique_name)
        # file.save(save_path) 

                # 保存文件到本地
        with open(save_path, "wb") as buffer:
            content = await file.read() 
            buffer.write(content) 

        # 核心修正点：正确的数据解构
        result = ocr_engine.ocr(save_path,  cls=True)
        texts = [line[1][0] for res in result for line in res]

        # 删除文件
        os.remove(save_path)

        # 删除tmp.jpg
        if os.path.exists('tmp.jpg'):
            os.remove('tmp.jpg')

        # 返回结果
        return JSONResponse(texts)
        
    except Exception as e:
        return JSONResponse({"error": str(e)}), 500 